# -*- coding: utf-8 -*-
# Author   : ZhangQing
# Time     : 2025-08-12 10:45
# File     : security_middleware.py
# Project  : codebuddy_craft
# Desc     : API安全中间件

"""
API安全中间件 🛡️

提供API安全防护，包括请求ID生成、速率限制、安全头设置等
"""

import time
import uuid
import logging
from typing import Callable, Dict, Optional, Set, List, Any

from fastapi import FastAPI, Request, Response
from fastapi.middleware.cors import CORSMiddleware
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.types import ASGIApp

from backend.core.settings import get_settings
from backend.utils.logger import log_api_request, log_api_response

settings = get_settings()
logger = logging.getLogger(__name__)


class RequestIDMiddleware(BaseHTTPMiddleware):
    """请求ID中间件 🔍
    
    为每个请求生成唯一ID，用于跟踪和日志记录
    """
    
    def __init__(self, app: ASGIApp, header_name: str = "X-Request-ID"):
        super().__init__(app)
        self.header_name = header_name
    
    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        # 从请求头获取请求ID，如果没有则生成新的
        request_id = request.headers.get(self.header_name) or str(uuid.uuid4())
        
        # 将请求ID存储在请求状态中
        request.state.request_id = request_id
        
        # 处理请求
        response = await call_next(request)
        
        # 在响应头中添加请求ID
        response.headers[self.header_name] = request_id
        
        return response


class SecurityHeadersMiddleware(BaseHTTPMiddleware):
    """安全头中间件 🔒
    
    添加安全相关的HTTP头
    """
    
    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        response = await call_next(request)
        
        # 添加安全头
        response.headers["X-Content-Type-Options"] = "nosniff"
        response.headers["X-Frame-Options"] = "DENY"
        response.headers["X-XSS-Protection"] = "1; mode=block"
        response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
        
        # 内容安全策略
        if settings.is_production:
            csp_directives = [
                "default-src 'self'",
                "img-src 'self' data:",
                "style-src 'self' 'unsafe-inline'",
                "script-src 'self'",
                "connect-src 'self'",
                "font-src 'self'",
                "object-src 'none'",
                "base-uri 'self'",
                "form-action 'self'",
                "frame-ancestors 'none'"
            ]
            response.headers["Content-Security-Policy"] = "; ".join(csp_directives)
        
        # 严格传输安全
        if settings.is_production:
            response.headers["Strict-Transport-Security"] = "max-age=31536000; includeSubDomains"
        
        return response


class RateLimitMiddleware(BaseHTTPMiddleware):
    """速率限制中间件 ⏱️
    
    限制API请求频率，防止滥用
    """
    
    def __init__(
        self, 
        app: ASGIApp, 
        rate_limit: int = 60,  # 每分钟请求数
        window_size: int = 60,  # 窗口大小（秒）
        exclude_paths: Optional[Set[str]] = None  # 排除的路径
    ):
        super().__init__(app)
        self.rate_limit = rate_limit
        self.window_size = window_size
        self.exclude_paths = exclude_paths or set()
        self.requests: Dict[str, List[float]] = {}  # IP -> 请求时间戳列表
    
    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        # 检查是否排除当前路径
        if request.url.path in self.exclude_paths:
            return await call_next(request)
        
        # 获取客户端IP
        client_ip = request.client.host if request.client else "unknown"
        
        # 当前时间
        current_time = time.time()
        
        # 清理过期的请求记录
        if client_ip in self.requests:
            self.requests[client_ip] = [
                ts for ts in self.requests[client_ip]
                if current_time - ts < self.window_size
            ]
        else:
            self.requests[client_ip] = []
        
        # 检查是否超过速率限制
        if len(self.requests[client_ip]) >= self.rate_limit:
            logger.warning(
                f"速率限制: IP {client_ip} 超过限制 ({self.rate_limit}/{self.window_size}s)",
                extra={
                    "client_ip": client_ip,
                    "path": request.url.path,
                    "method": request.method,
                    "rate_limit": self.rate_limit,
                    "window_size": self.window_size
                }
            )
            
            # 返回429错误
            from backend.api.error_handlers import format_error_response
            return Response(
                content=format_error_response(
                    status_code=429,
                    message=f"请求过于频繁，请在{self.window_size}秒后重试",
                    error_code="rate_limit_exceeded"
                ),
                status_code=429,
                media_type="application/json"
            )
        
        # 记录请求时间
        self.requests[client_ip].append(current_time)
        
        # 处理请求
        return await call_next(request)


class APILoggingMiddleware(BaseHTTPMiddleware):
    """API日志中间件 📝
    
    记录API请求和响应日志
    """
    
    def __init__(
        self, 
        app: ASGIApp, 
        exclude_paths: Optional[Set[str]] = None,  # 排除的路径
        exclude_methods: Optional[Set[str]] = None  # 排除的方法
    ):
        super().__init__(app)
        self.exclude_paths = exclude_paths or {"/health", "/metrics"}
        self.exclude_methods = exclude_methods or {"OPTIONS"}
    
    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        # 检查是否排除当前请求
        if request.url.path in self.exclude_paths or request.method in self.exclude_methods:
            return await call_next(request)
        
        # 获取请求信息
        request_info = {
            "method": request.method,
            "path": request.url.path,
            "query_params": str(request.query_params),
            "client_ip": request.client.host if request.client else "unknown",
            "user_agent": request.headers.get("user-agent", "unknown"),
            "request_id": getattr(request.state, "request_id", None)
        }
        
        # 记录请求日志
        log_api_request(request_info)
        
        # 记录请求开始时间
        start_time = time.time()
        
        # 处理请求
        response = await call_next(request)
        
        # 计算请求处理时间
        duration = time.time() - start_time
        
        # 记录响应日志
        response_info = {
            "status_code": response.status_code,
            "duration": duration,
            "path": request.url.path,
            "method": request.method,
            "client_ip": request_info["client_ip"],
            "request_id": request_info["request_id"]
        }
        log_api_response(response_info)
        
        return response


def setup_security_middleware(app: FastAPI) -> None:
    """设置安全中间件 🛡️

    Args:
        app: FastAPI应用实例
    """
    # 添加CORS中间件
    app.add_middleware(
        CORSMiddleware,
        allow_origins=settings.CORS_ORIGINS,
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )
    
    # 添加请求ID中间件
    app.add_middleware(RequestIDMiddleware)
    
    # 添加安全头中间件
    app.add_middleware(SecurityHeadersMiddleware)
    
    # 添加API日志中间件
    app.add_middleware(
        APILoggingMiddleware,
        exclude_paths={"/health", "/metrics", "/docs", "/redoc", "/openapi.json"}
    )
    
    # 添加速率限制中间件（仅在生产环境）
    if settings.is_production:
        app.add_middleware(
            RateLimitMiddleware,
            rate_limit=settings.RATE_LIMIT,
            window_size=60,
            exclude_paths={"/health", "/metrics", "/docs", "/redoc", "/openapi.json"}
        )


# 导出所有公共类和函数
__all__ = [
    'RequestIDMiddleware',
    'SecurityHeadersMiddleware',
    'RateLimitMiddleware',
    'APILoggingMiddleware',
    'setup_security_middleware',
]